# -*- coding: utf-8 -*-
"""
Created on Tue Jul  6 11:15:49 2021

@author: sarupa
"""
# Define symbolic variable and model
from __future__ import print_function, division # Grab some handy Python3 stuff.
import mpctools as mpc
import numpy as np
import casadi
import control
from casadi import *
import matplotlib.pyplot as plt
from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler

# Number of variables and order
Nx = 9 # Number of states
Nu = 7 # Number of inputs
Ny = 1 # Number of target output
order = Nx # Number of Lie derivative from 0 to order-1

#define variables
F10=1.50406459e+01
F20= 3.42104461e+00
Q1=3.14595480e+06
Q2= 1.02229466e+06
Q3=  3.49448744e+06 
Fr= 1.50551530e+01 
Fp=6.52273829e-01

T10=300
T20=300
V1= 4.0
V2= 4.0
V3= 4.0

E1= 50000
E2= 60000
k1= 9.972e06
k2= 9.36e06

dH1= -6*10000
dH2= -7*10000
aA= 3.5
aB= 1
aC= 0.5
Cp= 4.2
R= 8.314
rho= 1000
xA10=1.0
xA20=1.0
xB10= 0.0
xB20= 0.0
Hvap1= -3.57*10000
Hvap2= -1.57*10000
Hvap3= -4.07*10000
cm= 2
       
#system disturance
w1=0.0
w2=0.0
w3=0.0
w4=0.0
w5=0.0
w6=0.0
w7=0.0
w8=0.0
w9=0.0
# symbolic varaible
x = casadi.SX.sym('x',Nx,1) 
u = casadi.SX.sym('x',Nu,1) 

# symbolic ode model f
f = casadi.SX.sym('f',Nx,1)

xA1 = x[0]
xB1 = x[1]
T1 = x[2]
xA2 = x[3]
xB2 = x[4]
T2 = x[5]
xA3 = x[6]
xB3 = x[7]
T3 = x[8]

xC3 = 1 - xA3 - xB3
x3a = aA*xA3 + aB*xB3 + aC*xC3
xAr = aA*xA3/x3a
xBr = aB*xB3/x3a
xCr = aC*xC3/x3a

F1 = F10 
F2 =F1 + F20

f[0]=F10*(xA10 - xA1)/V1 - k1*exp(-E1/(R*T1))*xA1 +w1
f[1]=F10*(xB10 - xB1)/V1 + k1*exp(-E1/(R*T1))*xA1 - k2*exp(-E2/(R*T1))*xB1 +w2
f[2]=F10*(T10 - T1)/V1- dH1*k1*exp(-E1/(R*T1))*xA1/Cp*cm/rho - dH2*k2*exp(-E2/(R*T1))*xB1/Cp*cm/rho + Q1/(rho*Cp*V1) +w3
f[3]=F1*(xA1 - xA2)/V2 + F20*(xA20 - xA2)/V2 - k1*exp(-E1/(R*T2))*xA2 +w4
f[4]=F1*(xB1 - xB2)/V2 + F20*(xB20 - xB2)/V2 + k1*exp(-E1/(R*T2))*xA2 - k2*exp(-E2/(R*T2))*xB2 +w5
f[5]=F1*(T1 - T2)/V2 + F20*(T20-T2)/V2 - dH1*k1*exp(-E1/(R*T2))*xA2/Cp *cm/rho - dH2*k2*exp(-E2/(R*T2))*xB2/Cp *cm/rho+ Q2/(rho*Cp*V2) +w6
f[6]=F2*(xA2 - xA3)/V3 - (Fr+Fp)*(xAr - xA3)/V3 +w7
f[7]=F2*(xB2 - xB3)/V3 - (Fr+Fp)*(xBr - xB3)/V3 +w8
f[8]=F2*(T2 - T3)/V3 + Q3/(rho*Cp*V3)+(Fr+Fp)*(xAr*Hvap1+xBr*Hvap2+xCr*Hvap3)/(rho*Cp*V3)*cm +w9

g = casadi.SX.sym('g',Nu,1) # u_dot = g(x,u); For constant u, g = 0.
g[0,0] = 0
g[1,0] = 0
g[2,0] = 0
g[3,0] = 0
g[4,0] = 0
g[5,0] = 0
g[6,0] = 0

# Measurements
h = casadi.SX.sym('h',Ny,1) # output function
h[0,0] = x[7]

print('Symbolic model is created.')

# Code : Calculate Y_Lie and O_Lie: Numerical O_Lie at steady-state (xs,us)

xval =np.array([9.06393325e-01, 9.32535771e-02, 3.52482083e+02, 5.80306393e-01,
       4.08395110e-01, 3.98003376e+02, 3.18882689e-01, 6.51873344e-01,
       3.99773315e+02]).T
xs = xval

x0=np.array([.25, .55, 460, .25, .6, 475, .1, 0.55, 480])*0.95

uval=np.array([1.50406459e+01, 3.42104461e+00, 3.14595480e+06, 1.50551530e+01,
       3.49448744e+06, 1.02229466e+06, 6.52273829e-01])*1.03

# Linearized model and observability matrix
A_sym = casadi.SX.sym('A_sym',Nx,Nx)
A_sym = casadi.jacobian(f,x)
A_fun = casadi.Function('A_fun',[x,u],[A_sym[:,:]])
a = A_fun(xval,uval) #Linearized A
As = np.array(a)

C_sym = casadi.SX.sym('C_sym',Ny,Nx)
C_sym = casadi.jacobian(h,x)
C_fun = casadi.Function('C_fun',[x,u],[C_sym[:,:]])
a = C_fun(xval,uval) # Linearized A
Cs = np.array(a)

O_Lin_ss = control.obsv(As,Cs)
Rank_O_Lin_ss = np.linalg.matrix_rank(O_Lin_ss)
print('O_Lin_ss finished.')

# Code : Define the CSTR model
Nsim = 200 # Number of sampling time
Nw = Nx # Number of process noise
Nv = Ny # Number of measurement noise
Delta = 0.01 # Sampling time
    
def cstrmodel(x,u):
    
    xA1 = x[0]
    xB1 = x[1]
    T1 = x[2]
    xA2 = x[3]
    xB2 = x[4]
    T2 = x[5]
    xA3 = x[6]
    xB3 = x[7]
    T3 = x[8]
        
    xC3 = 1 - xA3 - xB3
    x3a = aA*xA3 + aB*xB3 + aC*xC3
    xAr = aA*xA3/x3a
    xBr = aB*xB3/x3a
    xCr = aC*xC3/x3a
    
    F10= u[0]
    F20= u[1]
    Q1= u[2]
    Fr= u[3]
    Q3= u[4]
    Q2= u[5]
    Fp= u[6]
    
    F1 = F10 
    F2 =F1 + F20
    
    dxdt= [ F10*(xA10 - xA1)/V1 - k1*exp(-E1/(R*T1))*xA1 +w1,
         F10*(xB10 - xB1)/V1 + k1*exp(-E1/(R*T1))*xA1 - k2*exp(-E2/(R*T1))*xB1 +w2,
         F10*(T10 - T1)/V1- dH1*k1*exp(-E1/(R*T1))*xA1/Cp*cm/rho - dH2*k2*exp(-E2/(R*T1))*xB1/Cp*cm/rho + Q1/(rho*Cp*V1) +w3,
         F1*(xA1 - xA2)/V2 + F20*(xA20 - xA2)/V2 - k1*exp(-E1/(R*T2))*xA2 +w4,
         F1*(xB1 - xB2)/V2 + F20*(xB20 - xB2)/V2 + k1*exp(-E1/(R*T2))*xA2 - k2*exp(-E2/(R*T2))*xB2 +w5,
         F1*(T1 - T2)/V2 + F20*(T20-T2)/V2 - dH1*k1*exp(-E1/(R*T2))*xA2/Cp *cm/rho - dH2*k2*exp(-E2/(R*T2))*xB2/Cp *cm/rho+ Q2/(rho*Cp*V2) +w6, 
         F2*(xA2 - xA3)/V3 - (Fr+Fp)*(xAr - xA3)/V3 +w7,
         F2*(xB2 - xB3)/V3 - (Fr+Fp)*(xBr - xB3)/V3 +w8,
         F2*(T2 - T3)/V3 + Q3/(rho*Cp*V3)+(Fr+Fp)*(xAr*Hvap1+xBr*Hvap2+xCr*Hvap3)/(rho*Cp*V3)*cm +w9
         ]
    return dxdt

F = mpc.getCasadiFunc(cstrmodel,[Nx,Nu],["x","u"],"cstrmodel")
F_rk4 = mpc.getCasadiFunc(cstrmodel,[Nx,Nu],["x","u"],"ode_rk4",rk4=True,Delta=Delta,M=1)

cstr = mpc.DiscreteSimulator(cstrmodel, Delta, [Nx,Nu], ["x","u"])

# measurement
def measurement(x):
    return x[7]
ys = measurement(xs)

# Turn into casadi functions.
H = mpc.getCasadiFunc(measurement,[Nx],["x"],"measurement")

print('Model is created.')

# Code : Generate test data including u(0:Nsim-1) and y(0:Nsim)
## Generate test data including u(0:Nsim-1) and y(0:Nsim)

#'''
# generate data w & v
sigma_w = 0.001#0.05#0.01
sigma_v = 0.001#0.05#0.01
v = sigma_v*np.random.randn(Nsim,Nv)
usim = np.zeros((Nsim+1,Nu))
usim[:,:] = uval
np.random.seed(100)
w=sigma_w*np.random.randn(Nsim,Nw)
w[:,2]=.100*w[:,2]
w[:,5]=.100*w[:,5]
w[:,8]=.100*w[:,8]

# generate data x and y
xsim = np.zeros((Nsim+1,Nx))
xsim[0,:] = x0
yclean = np.zeros((Nsim, Ny))
ysim = np.zeros((Nsim, Ny))

for t in range(Nsim):
    yclean[t,:] = measurement(xsim[t,:]) # Get zero-noise measurement.
    ysim[t,:] = yclean[t,:] + v[t,:] # Add noise to measurement.
    xsim[t+1,:] = cstr.sim(xsim[t,:],usim[t,:]) + w[t,:]

print('Data is generated.')


# Code 7: Create the function including sensitivity analysis and variable selection
# Create the function including sensitivity analysis and variable selection

def SensAnal(size,x,y,u):
    print(size)
    Sxp = np.zeros([size+1,Nx,Nx])
    Sxp[0,:,:] = np.eye(Nx)
    Syp = np.zeros([size+1,Ny,Nx])
    
    dFdx = np.zeros([size,Nx,Nx])
    dHdx = np.zeros([size+1,Ny,Nx])
        
    x_sen = np.zeros((size+1,Nx))
    y_sen = np.zeros((size+1,Ny))
    u_sen = np.zeros((size,Nu))
    
    x_sen = x
    y_sen = y
    u_sen = u
    
    for i in range(size):
        # Sensitivity matrix calculation
        JF = mpc.util.getLinearizedModel(F_rk4, [x_sen[i,:],u_sen[i,:]],
                                         ["A","B"],None)
        dFdx[i,:,:] = JF["A"]
        
        JH = mpc.util.getLinearizedModel(H, [x_sen[i,:]],["A"],None)
        dHdx[i,:,:] = JH["A"]
    
        Sxp[i+1,:,:] = mpc.mtimes(dFdx[i,:,:],Sxp[i,:,:])
        Syp[i,:,:] = mpc.mtimes(dHdx[i,:,:],Sxp[i,:,:])
    
    if size == 0:
        JH = mpc.util.getLinearizedModel(H, [x_sen[0,:]], ["A"],None)
        dHdx[0,:,:] = JH["A"]
        Syp[0,:,:] = mpc.mtimes(dHdx[0,:,:],Sxp[0,:,:])
    else:
        JH = mpc.util.getLinearizedModel(H, [x_sen[i+1,:]], ["A"],None)
        dHdx[i+1,:,:] = JH["A"]
        Syp[i+1,:,:] = mpc.mtimes(dHdx[i+1,:,:],Sxp[i+1,:,:])

    
    # Original sensitivity matrix
    for i in range(size+1):
        # Concatenate sensitivity matrixes
        if i==0:
            Sall = Syp[i,:,:]
        else:
            Sall = np.concatenate((Sall,Syp[i,:,:]), axis=0)
    # Normalisation
    for i in range(size+1):
        for i1 in range(Ny):
            for i2 in range(Nx):
                Syp[i,i1,i2] = Syp[i,i1,i2]*x_sen[0,i2]/y_sen[i,i1]
        # Concatenate sensitivity matrixes
        if i==0:
            Sall = Syp[i,:,:]
        else:
            Sall = np.concatenate((Sall,Syp[i,:,:]), axis=0)
            
    rank = np.linalg.matrix_rank(Sall)
    rank_S = rank
    cond = np.linalg.cond(Sall)

    U_svd, s_svd, Vh_svd = np.linalg.svd(Sall)
    
    sigma_log=np.log10(s_svd)
    
    return sigma_log, rank, U_svd, s_svd, Vh_svd 

print('The function of sensitivity evaluation and variable selection is created.')

# Code : Evaluate the ranks of O_Lie, O_Lin, and O_Sen along the trajectory

# In[8]:

# Evaluate the ranks of O_Lie, O_Lin, and O_Sen along the trajectory
Rank_O_Lie_tra = np.zeros(Nsim)
Rank_O_Lin_tra = np.zeros(Nsim)
Rank_O_Sen_tra = np.zeros(Nsim)

for i1 in range(Nsim-1,Nsim):
    print('\nSampling time: ', i1)
    
    # Evaluate the O_Lin
    b = A_fun(xsim[i1,:],usim[i1,:]) # Linearized A
    As = np.array(b)
    
    c = C_fun(xsim[i1,:],usim[i1,:]) # Linearized A
    Cs = np.array(c)
    
    O_Lin_tra = control.obsv(As,Cs)
    
    # Rank of O_Lin
    Rank_O_Lin_tra[i1] = np.linalg.matrix_rank(O_Lin_tra)
    print('Rank of O_Lin: ', Rank_O_Lin_tra[i1])
    
    # Rank of O_Sen
    Flag,rank, U_svd, s_svd, Vh_svd  = SensAnal(i1,xsim[0:i1+1,:],ysim[0:i1+1,:],usim[0:i1,:])
    sigma_log=Flag[0:9]
   
    print('Rank of O_Sen: ',rank)
    
t=np.arange(1, 10)
plt.figure()
plt.rcParams["figure.figsize"] = [5.00,4.0]
plt.rcParams["figure.autolayout"] = True
plt.plot(t,sigma_log,'o',linewidth=8) 
plt.axhline(y = 0.43 ,color = 'r', linestyle = '--')
plt.axhline(y= -0.12, color='r',linestyle ='--')  
plt.xlabel('No of singular components')
plt.ylabel('$Log_{10}$ of Singular Values')
plt.title('Explained variance: Elbow graph')
plt.show()
plt.savefig("Elbow_plot_state.jpg", dpi=200)

# sudden jump observed first between 3 and 4 singular value
diff=np.zeros(8)
for i in range(8):
    diff[i]=abs(sigma_log[i]-sigma_log[i+1])
print('diff=', diff)       

###The D_j is calculated and reported based on different initial condition, input values, steady state, 
#  different system disturbance and measurement noise

# This is a sample calculation
D_j=np.zeros(9)  
D_j= (np.abs(Vh_svd[0,:]*s_svd[0])+np.abs(Vh_svd[1,:]*s_svd[1])+np.abs(Vh_svd[2,:]*s_svd[2]))/sum(np.abs(s_svd[0:3]))
sum_vsort=np.argsort(D_j)
sum_sort=np.sort(D_j)
print("increasing order=",sum_sort)
print("increasing order=",sum_vsort)